import numpy as np
import laspy
import networkx as nx
import matplotlib.pyplot as plt
import open3d as o3d
import geojson
import time
import sys
import os

import xml.etree.ElementTree as ET
from sklearn.neighbors import KDTree
from configparser import ConfigParser
from scipy.spatial import Delaunay, ConvexHull,Voronoi
from scipy.interpolate import splprep, splev
from shapely.geometry import Polygon,LineString
from tqdm import tqdm
from joblib import Parallel, delayed

class PointCloudHandler:
    def __init__(self, point_cloud, traj_coords):
        """
        初始化点云处理器。

        参数:
        point_cloud (numpy.ndarray): 点云数据，形状为 (n, 3)，每行是 [x, y, z]。
        """
        self.point_cloud = point_cloud
        self.traj_coords = traj_coords


    def get_projected_z(self, traj_coords, point):
        """
        计算点到轨迹的投影点的 z 值。

        参数:
        traj_coords (numpy.ndarray): 轨迹坐标数组，形状为 (n, 3)。
        point (numpy.ndarray): 查询点，形状为 (3,)。

        返回:
        float: 投影点的 z 值。
        """
        traj_xy = traj_coords[:, :2]
        traj_z = traj_coords[:, 2]

        # 找到轨迹上最近的点
        distances = np.linalg.norm(traj_xy - point[:2], axis=1)
        nearest_index = np.argmin(distances)

        # 返回最近点的 z 值
        return traj_z[nearest_index]

    def get_road_boundary_index(self, sorted_distances, toleracence):
        i = 0
        while(i < len(sorted_distances) - 1):
            if (abs(sorted_distances[i+1] - sorted_distances[i]) <= toleracence):
                i += 1
            else:
                return i
        return i
    
    def voxel_filter(self, point_cloud, voxel_size):
        """
        对点云数据进行体素滤波。

        参数:
        voxel_size (float): 体素的大小。

        返回:
        numpy.ndarray: 经过体素滤波后的点云数据。
        """
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(point_cloud)
        pcd_down = pcd.voxel_down_sample(voxel_size=voxel_size)
        return np.asarray(pcd_down.points)

    def generate_boundary_lines(self, default_span = 5.0):
        """
        根据轨迹生成左右边界折线段。

        参数:
        step (int): 轨迹采样步长，默认 5。
        width (float): 路面宽度的一半，默认 5 米。

        返回:
        tuple: 左边界点列表和右边界点列表。
        """
        left_boundary = []
        right_boundary = []
        traj_coords = np.array(self.traj_coords)
        road_kdtree = KDTree(self.point_cloud[:, :2])  # 只使用 x 和 y 坐标构建 KDTree

        rdistances = []
        ldistances = []
        traj_points = []
        perpendiculars = []
        for i in tqdm(range(0, len(traj_coords) - 1)):
            # 当前轨迹点和下一个轨迹点
            p1 = traj_coords[i]
            p2 = traj_coords[i + 1]

            direction = p2[:2] - p1[:2]
            direction /= np.linalg.norm(direction)  # 单位化方向向量
            perpendicular = np.array([-direction[1], direction[0]])  # 垂直方向向量

            indices_list = road_kdtree.query_radius(p1[:2].reshape(1, -1), r=10)
            indices_list = np.concatenate(indices_list)

            if len(indices_list) != 0:
                points = self.point_cloud[indices_list]
                vdis = np.cross(points[:, :2] - p1[:2], perpendicular)
                condition = np.logical_and(
                    abs(points[:, 2] - p1[2]) < 0.1,
                    abs(vdis) < 0.5
                )
                condition_indices = np.where(condition)[0]
                condition_indices = indices_list[condition_indices]
                boundary_points = self.point_cloud[condition_indices]

                if len(boundary_points) != 0:
                    distances = np.dot(boundary_points[:, :2] - p1[:2], perpendicular)
                    positive_distances = distances[distances > 0]
                    negative_distances = distances[distances < 0]
                    sorted_positive = np.sort(positive_distances)  # 升序排列正距离
                    sorted_negative = np.sort(negative_distances)[::-1]  # 降序排列负距离(绝对值从小到大)

                    #计算左右边界点
                    ld = default_span
                    if len(sorted_positive) > 0:
                        idx = self.get_road_boundary_index(sorted_positive, 0.2)
                        distance = abs(sorted_positive[idx])
                        if distance > 2.0 and distance < 10.0:
                            ld = distance
                    rd = default_span
                    if len(sorted_negative) > 0:
                        idx = self.get_road_boundary_index(sorted_negative, 0.2)
                        distance = abs(idx)
                        if distance > 2.0 and distance < 10.0:
                            rd = distance
                    ldistances.append(ld)
                    rdistances.append(rd)
                    perpendiculars.append(perpendicular)
                    traj_points.append(p1)
        ldistances = self.smooth(ldistances)
        rdistances = self.smooth(rdistances)
        for i in range(len(traj_points)):
            left_point = traj_points[i][:2] - perpendiculars[i] * ldistances[i]
            right_point = traj_points[i][:2] + perpendiculars[i] * rdistances[i]
            left_boundary.append([left_point[0], left_point[1], traj_points[i][2]])
            right_boundary.append([right_point[0], right_point[1], traj_points[i][2]])
        return left_boundary, right_boundary
    
    def smooth(self, distances, span = 30):
        if len(distances) < span:
            return distances
        
        new_distances = distances.copy()
        for i in range(0, len(distances)):
            start_index = i - span
            end_index = i
            if i < span:
                start_index = 0
                end_index = span

            if i + span > len(distances):
                start_index = len(distances) - span
                end_index = len(distances)
            
            new_distances[i] = np.mean(distances[start_index:end_index])
        return new_distances

        
    def get_z(self, x, y):
        """
        获取给定 (x, y) 对应的 z 值。

        参数:
        x (float): 查询点的 x 坐标。
        y (float): 查询点的 y 坐标。

        返回:
        float: 最近点的 z 值。
        """
        dist, index = self.kdtree.query([[x, y]], k=1)  # 查询最近的 1 个点
        return self.point_cloud[index[0][0], 2]
    

    def dbscan_optimized(self, data, eps, min_samples, n_jobs=-1):
        """
        优化后的 DBSCAN 实现，支持并行化。

        参数:
        data (numpy.ndarray): 输入数据点集。
        eps (float): 邻域半径。
        min_samples (int): 核心点的最少邻居数量。
        n_jobs (int): 并行线程数，默认为 -1（使用所有可用线程）。

        返回:
        numpy.ndarray: 聚类标签。
        """
        n = len(data)
        labels = np.zeros(n, dtype=int)  # 初始化标签，0 表示未分类，-1 表示噪声点
        cluster_id = 0

        kdt = KDTree(data, leaf_size=40)  # 构建 KDTree

        # 查询所有点的邻居
        neighbors_list = Parallel(n_jobs=n_jobs)(
            delayed(kdt.query_radius)(data[i].reshape(1, -1), r=eps, return_distance=False)
            for i in range(n)
        )

        for i in range(n):
            if labels[i] != 0:
                continue

            neighbors = neighbors_list[i][0].tolist()
            if len(neighbors) < min_samples:
                labels[i] = -1  # 标记为噪声点
            else:
                cluster_id += 1
                labels[i] = cluster_id
                stack = set(neighbors)

                while stack:
                    neighbor_index = stack.pop()
                    if labels[neighbor_index] == -1:
                        labels[neighbor_index] = cluster_id
                    elif labels[neighbor_index] == 0:
                        labels[neighbor_index] = cluster_id
                        new_neighbors = neighbors_list[neighbor_index][0].tolist()
                        if len(new_neighbors) >= min_samples:
                            stack.update(new_neighbors)

        return labels

    def alpha_shape(self, points, alpha):
        """
        计算点集的alpha shape。

        参数:
        points (numpy.ndarray)：二维点集，形状为 (n, 2)，n为点的数量。
        alpha (float)：alpha值，用于控制形状的“松紧”程度。

        返回:
        edge_points (list)：构成alpha shape边界的边的端点索引列表。
        """
        if len(points) < 4:
            # 如果点数小于4，直接返回凸包（简单情况）
            hull = ConvexHull(points)
            return [(hull.vertices[i], hull.vertices[(i + 1) % len(hull.vertices)]) for i in range(len(hull.vertices))]

        tri = Delaunay(points)
        edge_points = []
        edges = set()

        # 遍历所有三角形
        for ia, ib, ic in tri.simplices:
            pa = points[ia]
            pb = points[ib]
            pc = points[ic]

            # 计算外接圆半径和圆心
            a = np.linalg.norm(pb - pc)
            b = np.linalg.norm(pc - pa)
            c = np.linalg.norm(pa - pb)
            s = (a + b + c) / 2
            area = np.sqrt(s * (s - a) * (s - b) * (s - c))
            if area == 0 :
                continue

            # 如果外接圆半径小于等于alpha，则将三角形的边加入到边集合中
            circum_radius = a * b * c / (4 * area)
            if circum_radius <= alpha:
                edges.add((ia, ib))
                edges.add((ib, ic))
                edges.add((ic, ia))

        
        edge_points = []
        # 处理边集合，去除重复的边（因为在不同三角形中可能重复出现）
        for e in edges:
            if (e[1], e[0]) not in edges:
                edge_points.append(e)
            
        # 创建一个无向图
        G = nx.Graph()

        # 添加边到图中
        G.add_edges_from(edge_points)

        # 找到所有的简单环（边界环）
        boundaries = list(nx.simple_cycles(G))

        # 计算每个环的长度并找到最长的环
        longest_boundary = max(boundaries, key=len)
        return longest_boundary
    
    def smooth_polygon_z_values(self, polygons, window_size):
        """
        平滑多边形的 z 值。

        参数:
        polygons (list): 包含多边形顶点的列表，每个多边形是一个点的列表。
        window_size (int): 平滑窗口大小，默认为 3。

        返回:
        list: 平滑后的多边形列表。
        """
        smoothed_polygons = []
        for polygon in polygons:
            smoothed_polygon = []
            z_values = [point[2] for point in polygon]  # 提取 z 值
            n = len(z_values)

            # 对 z 值进行平滑
            smoothed_z_values = []
            for i in range(n):
                # 计算窗口范围
                start = max(0, i - window_size // 2)
                end = min(n, i + window_size // 2 + 1)
                # 计算窗口内的平均值
                smoothed_z = sum(z_values[start:end]) / (end - start)
                smoothed_z_values.append(smoothed_z)

            # 更新平滑后的顶点
            for i, point in enumerate(polygon):
                smoothed_polygon.append([point[0], point[1], smoothed_z_values[i]])
            smoothed_polygons.append(smoothed_polygon)
        return smoothed_polygons

    def savePolygonToGeoJson(self, json_file, polygons):
        """
        将多边形保存为 GeoJSON 文件，支持三维坐标。

        参数:
        json_file (str): 输出的 GeoJSON 文件路径。
        polygons (list): 包含多边形顶点的列表，每个多边形是一个点的列表。
        """
        # 创建一个空的 FeatureCollection
        feature_collection = geojson.FeatureCollection([])

        # 遍历每个多边形，创建 Polygon 对象，并将其添加到 FeatureCollection 中
        for idx, polygon in enumerate(polygons):
            # 创建 Polygon 对象，直接使用三维坐标
            polygon_geojson = geojson.Polygon([polygon])

            # 创建一个 Feature，并为其添加属性
            feature = geojson.Feature(
                geometry=polygon_geojson,
                properties={"id": idx + 1, "name": f"Polygon {idx + 1}"}
            )

            # 将 Feature 添加到 FeatureCollection 中
            feature_collection.features.append(feature)

        # 将 FeatureCollection 保存为 GeoJSON 文件
        with open(json_file, 'w') as f:
            geojson.dump(feature_collection, f)

    def save_boundary_lines_to_geojson(self, geojson_file, left_boundary, right_boundary):
        """
        将左右边界线保存为 GeoJSON 文件。

        参数:
        geojson_file (str): 输出的 GeoJSON 文件路径。
        left_boundary (list): 左边界线的点列表。
        right_boundary (list): 右边界线的点列表。
        """
        features = []

        # 创建左边界线的 Feature
        left_line = LineString(left_boundary)
        left_feature = geojson.Feature(geometry=left_line, properties={"name": "Left Boundary"})
        features.append(left_feature)

        # 创建右边界线的 Feature
        right_line = LineString(right_boundary)
        right_feature = geojson.Feature(geometry=right_line, properties={"name": "Right Boundary"})
        features.append(right_feature)

        # 创建 FeatureCollection
        feature_collection = geojson.FeatureCollection(features)

        # 保存为 GeoJSON 文件
        with open(geojson_file, 'w') as f:
            geojson.dump(feature_collection, f)


    def save_traj_coords_to_obj(self, traj_coords, obj_file):
        """
        将轨迹坐标保存为 .obj 文件。

        参数:
        obj_file (str): 输出的 .obj 文件路径。
        """
        with open(obj_file, 'w') as f:
            for coord in traj_coords:
                f.write(f"v {coord[0]} {coord[1]} {coord[2]}\n")
            for i in range(len(traj_coords) - 1):
                f.write(f"l {i + 1} {i + 2}\n")

def log_time(start_time, message):
    """
    计算并打印从 start_time 到当前时间的耗时。

    参数:
    start_time (float): 起始时间（time.time() 的返回值）。
    message (str): 要打印的消息，描述当前操作。
    
    返回:
    float: 当前时间（作为下一次调用的起始时间）。
    """
    end_time = time.time()
    elapse_time = end_time - start_time
    print(f"{message} elapse {elapse_time:.2f} s")
    return end_time

def read_trajectory_coordinates(traj_file):
    if traj_file.endswith(".osm"):
        return read_xml_and_create_polyline(traj_file)
    elif traj_file.endswith(".geojson"):
        return read_geojson_and_create_polyline(traj_file)
    else:
        return None

def read_geojson_and_create_polyline(traj_file):
    with open(traj_file, 'r') as f:
        geojson_data = geojson.load(f)
    
    coordinates = []
    if 'features' in geojson_data and geojson_data['features']:
        coordinates = geojson_data['features'][0]['geometry']["coordinates"]
    return coordinates


def read_xml_and_create_polyline(traj_file):
    # 解析XML文件
    tree = ET.parse(traj_file)
    root = tree.getroot()

    # 提取utm_east和utm_north的值
    coordinates = []
    for node in root.findall('node'):
        utm_east = float(node.find("tag[@k='utm_east']").attrib['v'])
        utm_north = float(node.find("tag[@k='utm_north']").attrib['v'])
        altitude = float(node.find("tag[@k='altitude']").attrib['v'])
        coordinates.append((utm_east, utm_north, altitude))
    return coordinates

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: computeb [pcd_file] [traj_file] [output_path]")
        exit(0)

    cloud_file = sys.argv[1]
    traj_file = sys.argv[2]
    output_path = sys.argv[3]
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    start_time = time.time()
    file_name, file_extension = os.path.splitext(cloud_file)
    if file_extension == ".pcd":
        pcd = o3d.io.read_point_cloud(cloud_file)
    elif file_extension == ".las":
        las = laspy.read(cloud_file)
        # 提取点云数据
        points = np.vstack((las.x, las.y, las.z)).T
        # 创建 Open3D 点云对象
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
    else:
        print(f"can not read cloud format {file_extension}")
        exit(0)

    config = ConfigParser()
    config.read('cfg.ini')
    voxel_size = float(config.get('param', 'voxel_size')) # 设置体素网格的大小，可以根据需要调整
    cluster_radius = float(config.get('param', 'cluster_radius'))  # 邻域半径，可根据实际数据调整
    flexibility  = float(config.get('param', 'flexibility'))
    min_samples = int(config.get('param', 'min_samples'))  # 核心点最少邻居数量，可调整
    start_time = log_time(start_time, "read cloud file")
    distance_threshold = float(config.get('param', 'distance_threshold'))  # 距离阈值
    height_threshold = float(config.get('param', 'height_threshold'))  # 高度阈值

    # 使用体素网格滤波法进行抽稀
    pcd = pcd.voxel_down_sample(voxel_size=voxel_size)
    centroid = np.asarray(pcd.points).mean(axis=0)
    points2d = np.asarray(pcd.points - centroid)[:,:2] 
    traj_coords = read_trajectory_coordinates(traj_file)
    if traj_coords is None:
        print(f"can not read trajectory file {traj_file}")
        exit(0)

    print(f"point size {len(pcd.points)}")
    print(f"traj coords size {len(traj_coords)}")

    # 执行DBSCAN聚类
    pcd_handler = PointCloudHandler(np.asarray(pcd.points), traj_coords)
    left_boundary, right_boundary = pcd_handler.generate_boundary_lines()
    start_time = log_time(start_time, "generate_boundary_lines")

    pcd_handler.save_boundary_lines_to_geojson(os.path.join(output_path, "bounding.geojson"), left_boundary, right_boundary)
    print(f"save boundary into {os.path.join(output_path, 'bounding.geojson')} done!")

    is_save_obj = True
    if is_save_obj:
        save_obj = os.path.join(output_path, "traj.obj")
        pcd_handler.save_traj_coords_to_obj(pcd_handler.traj_coords, save_obj)
        print(f"save skleton into {save_obj} done!")

        save_obj = os.path.join(output_path, "left_boundary.obj")
        pcd_handler.save_traj_coords_to_obj(left_boundary, save_obj)
        print(f"save skeleton into {save_obj} done!")

        save_obj = os.path.join(output_path, "right_boundary.obj")
        pcd_handler.save_traj_coords_to_obj(right_boundary, save_obj)
        print(f"save skeleton into {save_obj} done!")


